import networkx as nx
import pulp
import numpy as np
import matplotlib.pyplot as plt

# 管网拓扑生成（随机多源多泄漏）
def generate_gas_network(num_nodes=30, leak_nodes=3):
    G = nx.DiGraph()
    
    # 节点类型：0=气源，1=阀门，2=用户区，3=连接点
    node_types = {0: 'source', 1: 'valve', 2: 'user', 3: 'junction'}
    positions = {}
    
    # 添加节点属性
    for i in range(num_nodes):
        node_type = np.random.choice([0,1,2,3], p=[0.05, 0.3, 0.4, 0.25])
        G.add_node(i, 
                   type=node_types[node_type],
                   pressure=np.random.uniform(0.2, 1.0),  # MPa
                   diameter=np.random.choice([0.1, 0.2, 0.3, 0.4, 0.5]),  # m
                   users=np.random.randint(100,1000) if node_type==2 else 0,
                   leak=False)
        positions[i] = (np.random.rand(), np.random.rand())
    
    # 设置泄漏点
    leak_ids = np.random.choice(range(5, num_nodes), leak_nodes, replace=False)
    for i in leak_ids:
        G.nodes[i]['leak'] = True
    
    # 添加管道连接
    for i in range(num_nodes-1):
        for j in range(i+1, num_nodes):
            if np.random.rand() > 0.7:  # 随机连接
                # 管道属性
                length = np.sqrt((positions[i][0]-positions[j][0])**2 + 
                                (positions[i][1]-positions[j][1])**2)
                G.add_edge(i, j, 
                           length=round(length,2),
                           has_valve=np.random.rand()>0.8)  # 20%管道有阀门
    
    # 添加超级汇点（连接所有泄漏点）
    super_sink = num_nodes
    G.add_node(super_sink, type='super_sink')
    positions[super_sink] = (0.5, 0.5)
    for node in leak_ids:
        G.add_edge(node, super_sink, capacity=float('inf'))
    
    return G, positions

# 最小割算法（隔离泄漏点）
def min_cut_isolation(G, source_nodes):
    # 设置边容量（基于管径和压力）
    for u, v in G.edges():
        if 'capacity' not in G[u][v]:
            d = G.nodes[u].get('diameter', 0.3)
            p = G.nodes[u].get('pressure', 0.5)
            G[u][v]['capacity'] = d**2 * p  # 流量容量因子
    
    # 找到所有气源节点
    sources = [n for n in G.nodes if G.nodes[n]['type']=='source']
    
    # 计算最小割（使用Boykov-Kolmogorov算法）
    super_sink = max(G.nodes)
    cut_value, partition = nx.minimum_cut(G, sources[0], super_sink, flow_func=nx.boykov_kolmogorov)
    
    # 提取需关闭的阀门
    valves_to_close = []
    for u, v in G.edges():
        if u in partition[0] and v in partition[1]:
            if G[u][v].get('has_valve', False):
                valves_to_close.append((u,v))
    
    return valves_to_close, cut_value

# 多目标线性规划优化
def optimize_valve_strategy(G, candidate_valves):
    # 创建问题实例
    prob = pulp.LpProblem("Valve_Optimization", pulp.LpMinimize)
    
    # 决策变量（是否关闭阀门）
    valve_vars = pulp.LpVariable.dicts("Valve", candidate_valves, cat='Binary')
    
    # 目标函数权重（可调整）
    w_user = 0.5  # 停气用户权重
    w_leak = 0.3  # 泄漏总量权重
    w_time = 0.2  # 关阀时间权重
    
    # 目标函数：最小化加权损失
    prob += w_user * user_impact(G, valve_vars) + \
            w_leak * leak_volume(G, valve_vars) + \
            w_time * valve_time(valve_vars)
    
    # 约束：必须隔离所有泄漏点
    for leak_node in [n for n in G.nodes if G.nodes[n]['leak']]:
        # 确保泄漏点被隔离（简化约束）
        prob += sum(valve_vars.get((u,v),0) for u,v in G.in_edges(leak_node)) >= 1
    
    # 求解
    prob.solve()
    
    # 提取最优解
    solution = [(u,v) for (u,v) in candidate_valves if valve_vars[(u,v)].value() > 0.5]
    return solution

# 辅助函数（目标函数计算）
def user_impact(G, valve_vars):
    """计算停气用户数（简化模型）"""
    # 实际应用需DFS/BFS遍历受影响区域
    return sum(valve_vars[(u,v)] * G.nodes[v].get('users',0) for (u,v) in valve_vars)

def leak_volume(G, valve_vars):
    """计算泄漏总量（kg）"""
    total_leak = 0
    for node in [n for n in G.nodes if G.nodes[n]['leak']]:
        p = G.nodes[node]['pressure']  # MPa
        d = G.nodes[node]['diameter']  # m
        area = np.pi * (d/2)**2  # 泄漏面积
        leak_rate = 0.6 * area * np.sqrt(2*p*1e6/0.8)  # kg/s
        total_leak += leak_rate * 30 * sum(  # 30s关闭时间
            valve_vars.get((u,node),0) for u in G.predecessors(node))
    return total_leak

def valve_time(valve_vars):
    """计算关阀总时间（假设每个阀门30秒）"""
    return 30 * sum(valve_vars[v] for v in valve_vars)

# 可视化
def visualize_network(G, pos, valves_to_close, leaks):
    plt.figure(figsize=(12,10))
    
    # 节点颜色编码
    node_color = []
    for n in G.nodes:
        if G.nodes[n]['type'] == 'source': node_color.append('gold')
        elif G.nodes[n]['type'] == 'valve': node_color.append('lightgreen')
        elif G.nodes[n]['leak']: node_color.append('red')
        else: node_color.append('skyblue')
    
    # 绘制节点和边
    nx.draw_networkx_nodes(G, pos, node_size=300, node_color=node_color)
    nx.draw_networkx_edges(G, pos, width=1.5, alpha=0.5)
    
    # 标记需关闭的阀门
    nx.draw_networkx_edges(G, pos, edgelist=valves_to_close, 
                          width=3, edge_color='red', style='dashed')
    
    # 添加标签
    labels = {n: f"{n}\n({G.nodes[n]['users']}户)" if G.nodes[n]['users']>0 else str(n) 
             for n in G.nodes}
    nx.draw_networkx_labels(G, pos, labels, font_size=9)
    
    plt.title("燃气管网泄漏事故关阀策略", fontsize=15)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig('gas_leak_response.png', dpi=300)
    plt.show()

# 主程序
if __name__ == "__main__":
    # 生成随机管网
    G, pos = generate_gas_network(num_nodes=30, leak_nodes=3)
    leaks = [n for n in G.nodes if G.nodes[n]['leak']]
    
    # 最小割初步隔离
    candidate_valves, cut_value = min_cut_isolation(G, source_nodes=[0])
    print(f"最小割发现 {len(candidate_valves)} 个候选阀门 (割值={cut_value:.2f})")
    
    # 多目标优化
    optimal_valves = optimize_valve_strategy(G, candidate_valves)
    print(f"优化选择关闭 {len(optimal_valves)} 个阀门")
    
    # 可视化结果
    visualize_network(G, pos, optimal_valves, leaks)
